package middleware

import (
	"github.com/gin-gonic/gin"
	bucket "github.com/juju/ratelimit"
	"go.uber.org/ratelimit"
	"net/http"
	"time"
)

const rateLimitNum = 1000 // 每秒请求次数

// RateLimitDrop 漏桶限流
// go.uber.org/ratelimit
// 漏桶是指我们有一个一直装满了水的桶，每过固定的一段时间即向外漏一滴水。
// 如果你接到了这滴水，那么你就可以继续服务请求，
// 如果没有接到，那么就需要等待下一滴水。
func RateLimitDrop() func(ctx *gin.Context) {
	limiter := ratelimit.New(rateLimitNum)
	return func(ctx *gin.Context) {
		now := time.Now()
		if t := limiter.Take().Sub(now); t > 0 {
			// 如果返回时间比当前时间要大，说明需要等待
			// 如果需要等待， time.sleep(t) 然后放行
			// 如果不需要等待请求时间， 直接abort
			ctx.JSON(http.StatusOK, gin.H{"errno": "-1", "errmsg": "接口已限流"})
			ctx.Abort()
			return
		}

		ctx.Next()
	}
}

// RateLimitBucket 令牌筒限流
// github.com/juju/ratelimit
// 令牌桶则是指匀速向桶中添加令牌，服务请求时需要从桶中获取令牌，令牌的数目可以按照需要消耗的资源进行相应的调整。
// 如果没有令牌，可以选择等待，或者放弃。
func RateLimitBucket() func(ctx *gin.Context) {
	limiter := bucket.NewBucket(time.Second, rateLimitNum)
	return func(ctx *gin.Context) {
		if limiter.TakeAvailable(1) != 1 {
			ctx.JSON(http.StatusOK, gin.H{"errno": "-1", "errmsg": "接口已限流"})
			ctx.Abort()
			return
		}

		// 如果取一个令牌能成功，就放行
		ctx.Next()
	}
}
